import ollama
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import FAISS
from langchain.llms import Ollama

# 初始化 Ollama 模型
ollama_model = Ollama(model='qwen3:1.7b')

# 初始化嵌入模型（用于将文本转换为向量）
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")

# 初始化 FAISS 向量存储（用于检索）
vectorstore = FAISS.from_texts(
    ["唐僧是取经的领导者，孙悟空是他的徒弟，猪八戒和沙僧也是他的徒弟。",
     "《西游记》讲述了唐僧师徒四人前往西天取经的故事。"],
    embeddings
)

def api_generate(text: str):
    print(f'提问：{text}')

    # 检索相关文档
    docs = vectorstore.similarity_search(text, k=1)
    context = "\n".join([doc.page_content for doc in docs])

    # 构造完整的提示，包含检索到的上下文
    full_prompt = f"{context}\n\n{ollama_model.template}{text}"

    # 调用 ollama.generate 方法，设置 stream=False
    response = ollama.generate(
        stream=False,  # 非流式
        model='qwen3:1.7b',  # 修改大模型名称
        prompt=full_prompt,
    )

    print('-----------------------------------------')
    print(response['response'])  # 输出完整的响应内容
    print('-----------------------------------------')
    print('总耗时：{}'.format(response['total_duration']))
    print('-----------------------------------------')

# 测试
api_generate(text='请说下西游记的主要人物，讲述了一个什么故事?')
